package com.neshai.webapp.controller;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.neshai.webapp.service.UsageService;
import com.neshai.webapp.utils.R;
import jakarta.annotation.PostConstruct;
import jakarta.servlet.http.HttpServletRequest;
import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.HttpStatusCode;
import org.springframework.http.MediaType;
import org.springframework.util.Assert;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.client.WebClientResponseException;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Map;

@Log4j2
@RestController
@RequestMapping({"/api"})
@RequiredArgsConstructor
public class ClaudeController {

    private static final String CLAUDE_API_URL = "https://api.anthropic.com/v1/messages";
    //private static final String CLAUDE_API_URL = "http://127.0.0.1:8080/api";
    @Value("${claude.apikey}")
    private String CLAUDE_API_KEY;

    private WebClient webClient;

    private final UsageService usageService;

    @PostConstruct
    public void init() {
        this.webClient = WebClient.builder()
                .defaultHeader(HttpHeaders.CONTENT_TYPE, "application/json")
                .build();
    }

    @PostMapping(value = "/claude/completions/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamCompletionsPost(@RequestBody Map<String, String> param, HttpServletRequest request) {
        String sessionId = param.get("sessionId");
        String prompt = param.get("prompt");
        String model = param.get("model");
        Assert.hasLength(sessionId, "sessionId is null");
        Assert.hasLength(prompt, "prompt is null");
        Assert.hasLength(model, "model is null");
        Object idObject = request.getAttribute("id");

        switch (model) {
            case "claude-3-5-sonnet-20240620", "claude-3-haiku-20240307":
                if (!usageService.usePremiumService(Integer.parseInt((String) idObject)))
                    return getErrorRes("reached the limit");
                break;
            case "claude-3-opus-20240229":
                if (!usageService.useBasicService(Integer.parseInt((String) idObject)))
                    return getErrorRes("reached the limit");
                break;
            default:
                return getErrorRes("error");
        }

        // 创建主 JSON 对象
        JSONObject jsonObject = new JSONObject();

        // 设置简单的键值对
        jsonObject.put("max_tokens", 1024);
        jsonObject.put("model", model);
        jsonObject.put("stream", true);

        // 创建 messages 数组
        JSONArray messagesArray = new JSONArray();

        // 创建 message 对象
        JSONObject messageObject = new JSONObject();
        messageObject.put("role", "user");
        messageObject.put("content", prompt);

        // 将 message 对象添加到 messages 数组中
        messagesArray.add(messageObject);

        // 将 messages 数组添加到主 JSON 对象中
        jsonObject.put("messages", messagesArray);

        return webClient.post()
                .uri(CLAUDE_API_URL)
                .header("x-api-key", CLAUDE_API_KEY)
                .header("anthropic-version", "2023-06-01")
                .bodyValue(jsonObject.toJSONString())
                .retrieve()
                .bodyToFlux(String.class)
                .onErrorResume(WebClientResponseException.class, ex -> {
                    HttpStatusCode statusCode = ex.getStatusCode();
                    HttpStatus status = HttpStatus.resolve(statusCode.value());
                    String res = ex.getResponseBodyAsString();
                    log.error("OpenAI API error: {} {}", status, res);
                    return Mono.error(new RuntimeException(res));
                });
    }

    private Flux<String> getErrorRes(String msg) {
        return Flux.create(emitter -> {
            emitter.next(" ");
            emitter.next(" ");
            emitter.next(JSON.toJSONString(R.fail(msg)));
            emitter.complete();
        });
    }
}

